#Initialises patches and genes
initialize<-function(Ni=10000){
  hapA<-rep(0:1, each=Ni)
  hapB<-rep(0, 2*Ni)
  patch<-rep(0:1, each=Ni)
  out<-cbind(hapA, hapB, patch)
  out
}

#selects individuals based on their fitness in a patch driven by hapA
  # s is strength of selection against non-local haplotypes (0<=s<=1)

selection<-function(pop, s, Ni, demog=F){
  if (demog==F){
    N0<-sum(pop[,"patch"]==0)
    N1<-sum(pop[,"patch"]==1)
  }
  else {
    N0<-Ni
    N1<-Ni
  }
  w<-pop[,"hapA"]==pop[,"patch"]
  w[w==0]<-1-s
  p0<-sample(which(pop[,"patch"]==0), size=N0, replace=T, prob=w[which(pop[,"patch"]==0)])
  p1<-sample(which(pop[,"patch"]==1), size=N1, replace=T, prob=w[which(pop[,"patch"]==1)])
  out<-pop[c(p0, p1), ]
  out
}

migration<-function(pop, m, Ni, bval){
  P0<-subset(pop, pop[,"patch"]==0) #matrix for each patch
  P1<-subset(pop, pop[,"patch"]==1)
  nhap<-table(pop[,"patch"], factor(pop[,"hapB"], levels=c(0,1))) #n haplotypes in each patch
  nP0B<-floor(m*nhap[1,1]) #n Bs to move
  nP1B<-floor(m*nhap[2,1])
  print(c(nP0B, nP1B))
  if (nP0B>0) nP0B<-sample(which(P0[,"hapB"]==0), nP0B, replace=F) #and their indices
  if (nP1B>0) nP1B<-sample(which(P1[,"hapB"]==0), nP1B, replace=F)
  P0[nP0B, "patch"]<-1
  P1[nP1B, "patch"]<-0

  nP0b<-floor((m+bval)*nhap[1,2]) #n bs to move
  nP1b<-floor((m+bval)*nhap[2,2])
  if (nP0b>0) nP0b<-sample(which(P0[,"hapB"]==1), nP0b, replace=F) #and their indices
  if (nP1b>0) nP1b<-sample(which(P1[,"hapB"]==1), nP1b, replace=F)
  P0[nP0b, "patch"]<-1
  P1[nP1b, "patch"]<-0
  
  rbind(P0, P1)
}


#recombines a fraction (r) of the gamete population
recombination<-function(pop, r){
  P0<-subset(pop, pop[,"patch"]==0) #matrix for each patch
  P1<-subset(pop, pop[,"patch"]==1)
  P0recomb<-sample(1:nrow(P0), floor(nrow(P0)*r))
  P1recomb<-sample(1:nrow(P1), floor(nrow(P1)*r))
  P0[P0recomb, "hapB"]<-sample(P0[P0recomb, "hapB"], length(P0[P0recomb, "hapB"]), replace=F)
  P1[P1recomb, "hapB"]<-sample(P1[P1recomb, "hapB"], length(P1[P1recomb, "hapB"]), replace=F)
  rbind(P0, P1)
}
  
plotter<-function(pop){
  phap<-table(pop[,"patch"], pop[,"hapB"])/nrow(pop) #prop haplotypes in each patch
  par(mfrow=c(2,1))
  barplot(height=phap[1,], ylim=c(0,1), legend.text=length(which(pop[,"patch"]==0)), main="Patch 1")
  barplot(height=phap[2,], ylim=c(0,1), legend.text=length(which(pop[,"patch"]==1)), main="Patch 2")
}
  
# runs population over time
  #bval is the modification to m
iterator<-function(Ni=10000, s, m, bval, init.gens, final.gens, demog=F){
  pop<-initialize(Ni)
  for (ii in 1:init.gens){
    #plotter(pop)
    pop<-recombination(pop, 0.5)
    pop<-selection(pop=pop, s=s, Ni=Ni, demog=demog)
    pop<-migration(pop=pop, m=m, Ni=Ni, bval=0)
  }
  pop[sample(x=1:(2*Ni), size=Ni, replace=F),"hapB"]<-1
  for (ii in 1:final.gens){
    #plotter(pop)
    phap<-table(pop[,"hapB"])
    if (sum(phap==0 | phap==2*Ni)>0) break
    pop<-recombination(pop, 0.5)
    pop<-selection(pop=pop, s=s, Ni=Ni, demog=demog)
    pop<-migration(pop=pop, m=m, Ni=Ni, bval=bval)
  }
    pop
}

# converts bvals to 0<m*<m.max given m 
logistic.m<-function(bval, m, m.max=0.5) {m.max/(1+exp(-bval))-m}
